import { types, visit } from "recast";

const { blockStatement } = types.builders;

const visitLoop = function (path) {
  if (path.node.body != null && path.node.body.type != "BlockStatement") {
    path.node.body = blockStatement([path.node.body]);
  }
  this.traverse(path);
};

/**
 * 找到所有的if和for语句，给他们增加花括号
 * @param {抽象语法树} ast
 */
const addBlock = (ast) => {
  visit(ast, {
    // 找到所有的if语句给他们增加花括号
    visitIfStatement: function (path) {
      if (
        path.node.consequent != null &&
        path.node.consequent.type != "BlockStatement"
      ) {
        path.node.consequent = blockStatement([path.node.consequent]);
      }
      if (
        path.node.alternate != null &&
        path.node.alternate.type != "BlockStatement"
      ) {
        path.node.alternate = blockStatement([path.node.alternate]);
      }
      this.traverse(path);
    },
    visitDoWhileStatement: visitLoop,
    // 找到所有的for循环,while循环给他们增加花括号
    visitForStatement: visitLoop,
    visitWhileStatement: visitLoop,
    visitForInStatement: visitLoop,
  });
};

export default addBlock;
